package com.dify.tool.common.utils;

import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch._types.query_dsl.Query;
import co.elastic.clients.elasticsearch.core.*;
import co.elastic.clients.elasticsearch.core.search.Hit;
import co.elastic.clients.elasticsearch.core.search.TotalHits;
import co.elastic.clients.elasticsearch.indices.CreateIndexRequest;
import co.elastic.clients.elasticsearch.indices.DeleteIndexRequest;
import co.elastic.clients.elasticsearch.indices.ExistsRequest;
import co.elastic.clients.json.JsonData;
import co.elastic.clients.transport.endpoints.BooleanResponse;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.function.Consumer;
import java.util.stream.Collectors;

/**
 * @author shl
 * @date 2025-07-28
 * @description
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class ElasticsearchUtil {

    private final ElasticsearchClient client;

    /**
     * 索引是否存在
     */
    public boolean indexExists(String indexName) {
        BooleanResponse exists = null;
        try {
            exists = client.indices()
                    .exists(ExistsRequest.of(e -> e.index(indexName)));
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        return exists.value();
    }

    /**
     * 创建索引
     */
    public void createIndex(String indexName) {
        try {
            if (!indexExists(indexName)) {
                client.indices().create(CreateIndexRequest.of(c -> c.index(indexName)));
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 删除索引
     */
    public void deleteIndex(String indexName) {
        try {
            if (indexExists(indexName)) {
                client.indices().delete(DeleteIndexRequest.of(d -> d.index(indexName)));
            }
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 索引单个文档
     */
    public <T> void indexDocument(String indexName, T document) {
        try {
            client.index(i -> i
                    .index(indexName)
                    .document(document)
            );
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 批量索引文档
     */
    public <T> void bulkIndexDocuments(String indexName, List<T> documents) {
        BulkRequest.Builder br = new BulkRequest.Builder();

        for (T document : documents) {
            br.operations(op -> op
                    .index(idx -> idx
                            .index(indexName)
                            .document(document)
                    )
            );
        }
        try {
            client.bulk(br.build());
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 通用搜索方法
     */
    public <T> SearchResult<T> search(String indexName,
                                      Class<T> documentType,
                                      Consumer<SearchRequest.Builder> searchConfigurer) {

        SearchResponse<T> response = null;
        try {
            response = client.search(s -> {
                searchConfigurer.accept(s);
                return s.index(indexName);
            }, documentType);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }

        TotalHits totalHits = response.hits().total();
        long total = totalHits != null ? totalHits.value() : 0;

        List<T> results = response.hits().hits().stream()
                .map(Hit::source)
                .collect(Collectors.toList());

        return new SearchResult<>(results, total);
    }

    /**
     * 根据ID获取文档
     */
    public <T> Optional<T> getById(String indexName, String id, Class<T> documentType) {
        GetResponse<T> response = null;
        try {
            response = client.get(g -> g
                            .index(indexName)
                            .id(id),
                    documentType);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }

        return Optional.ofNullable(response.source());
    }

    /**
     * 删除文档
     */
    public void deleteDocument(String indexName, String id) {
        try {
            client.delete(d -> d
                    .index(indexName)
                    .id(id)
            );
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 新增文档（自动生成ID）
     */
    public <T> String createDocument(String indexName, T document) {
        IndexResponse response = null;
        try {
            response = client.index(i -> i
                    .index(indexName)
                    .document(document)
            );
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        return response.id();
    }

    /**
     * 批量新增文档
     */
    public <T> void bulkCreateDocuments(String indexName, List<T> documents) {
        BulkRequest.Builder br = new BulkRequest.Builder();

        for (T document : documents) {
            br.operations(op -> op
                    .index(idx -> idx
                            .index(indexName)
                            .document(document)
                    )
            );
        }

        try {
            client.bulk(br.build());
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 异步新增文档
     */
    public <T> void createDocumentAsync(String indexName, T document) {
        CompletableFuture.runAsync(() -> {
            try {
                client.index(i -> i
                        .index(indexName)
                        .document(document)
                );
                log.debug("Async created document in {}", indexName);
            } catch (IOException e) {
                log.error("Async document creation failed", e);
            }
        });
    }

    /**
     * 更新文档
     */
    public <T> void updateDocument(String indexName, String id, T document) {
        try {
            client.update(u -> u
                            .index(indexName)
                            .id(id)
                            .doc(document),
                    document.getClass()
            );
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 条件更新属性值
     * <p>
     * indexName 索引名
     * sourceField 源属性名
     * sourceValue 源属性值
     * targetField 目标属性名
     * targetValue 目标属性值
     *
     * @return 更新的文档数量
     */
    public void updateByField(String indexName, String sourceField, Object sourceValue, String targetField, Object targetValue) {
        try {
            String source = "ctx._source.".concat(targetField).concat(" = params.").concat(String.valueOf(targetValue));
            // 1. 构建更新请求
            UpdateByQueryRequest request = UpdateByQueryRequest.of(u -> u
                    .index(indexName)                                                           // 指定索引名称
                    .query(Query.of(q -> q                                              // 构建查询条件
                            .term(t -> t                                               // 精确匹配 term 查询
                                    .field(sourceField)                                      // 查询字段
                                    .value(String.valueOf(sourceValue))                     // 查询值
                            )
                    ))
                    .script(s -> s                                                                          // 构建更新脚本
                            .inline(i -> i                                                                 // 内联脚本
                                    .source(source)                                                              // 更新脚本
                                    .params(String.valueOf(targetValue), JsonData.of(targetValue))              // 脚本参数
                            )
                    )
                    .refresh(true)                 // 操作完成后刷新索引
            );
            // 2. 执行更新操作
            client.updateByQuery(request);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }

    }

    /**
     * 搜索结果封装类
     */
    public record SearchResult<T>(List<T> results, long total) {
    }
}
